import rosbag
import rospy
from std_msgs.msg import Header
from sensor_msgs.msg import Imu
from cv_bridge import CvBridge
import cv2
import os

# 读取IMU文件
imu_data = []
with open('imu.txt', 'r') as imu_file:
    for line in imu_file:
        data = line.strip().split(' ')
        imu_data.append(data)

# 读取编码器文件
encoder_data = []
with open('encoder.txt', 'r') as encoder_file:
    for line in encoder_file:
        data = line.strip().split(' ')
        encoder_data.append(data)

# 创建ROS bag文件
bag = rosbag.Bag('imu_encoder_image.bag', 'w')

bridge = CvBridge()

# 读取并写入图片消息
for file in sorted(os.listdir()):
    if file.endswith('.jpg'):
        timestamp = int(file[:-4])
        image = cv2.imread(file)
        image_msg = bridge.cv2_to_imgmsg(image, 'bgr8')
        image_msg.header.stamp = rospy.Time.from_sec(timestamp / 1000000000.0)
        bag.write('/usb_cam/image_raw', image_msg, image_msg.header.stamp)

# 读取并写入IMU消息
for data in imu_data:
    timestamp = int(data[0])
    imu_msg = Imu()
    imu_msg.header.stamp = rospy.Time.from_sec(timestamp / 1000000000.0)
    imu_msg.angular_velocity.x = float(data[1])
    imu_msg.angular_velocity.y = float(data[2])
    imu_msg.angular_velocity.z = float(data[3])
    imu_msg.linear_acceleration.x = float(data[4])
    imu_msg.linear_acceleration.y = float(data[5])
    imu_msg.linear_acceleration.z = float(data[6])
    bag.write('/wit/imu', imu_msg, imu_msg.header.stamp)

# 读取并写入编码器数据消息
for data in encoder_data:
    timestamp = int(data[0])
    encoder_msg = Header()
    encoder_msg.stamp = rospy.Time.from_sec(timestamp / 1000000000.0)
    encoder_msg.frame_id = str(data[1])
    bag.write('/encoder_data', encoder_msg, encoder_msg.stamp)

# 关闭ROS bag文件
bag.close()